/*
 * Copyright 2012-2016 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package keter.sec.xss.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.ResponseBody;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.math.BigInteger;
import java.security.MessageDigest;

@Controller
public class CsrfTokenController {
  private static final Logger logger = LoggerFactory.getLogger(CsrfTokenController.class);

  /**
   * 利用跨域请求每次更换session，从而不包含session中的CSRF属性进行CSRF
   * @param request
   * @param response
   * @return
   */
  @GetMapping("/csrf")
  @ResponseBody
  public String csrf(HttpServletRequest request, HttpServletResponse response) {
    String s = getMD5(request.getSession().getId()); // 使用session生成MD5用做CSRF
    request.getSession().setAttribute("csrf_token", s);
    addAllowHeader(response);
    // 同域：request请求一次后sesson不变
    System.out.println("sid:[" + request.getSession().getId() + "] , token:" + s);
    return s;
  }

  /**
   * 利用跨域请求无法获取响应信息header中的csrf_coken数据
   * @param request
   * @param response
   * @return
   */
  @GetMapping("/csrfHeader")
  @ResponseBody
  public void csrfHeader(HttpServletRequest request, HttpServletResponse response) {
    response.setHeader("csrf_coken", getMD5(request.getSession().getId()));
    addAllowHeader(response);
  }

  @PostMapping("/csrf/form")
  @ResponseBody
  public String xsspost(HttpServletRequest request, HttpServletResponse response)
      throws ServletException, IOException {
    System.out.println("sid:" + request.getSession().getId());
    String csrf_token = request.getParameter("csrf_token");
    String csrf_token_in_session = request.getSession().getAttribute("csrf_token") + "";
    // 如果不添加允许跨域，已被浏览器阻止访问！
    addAllowHeader(response);
    // 跨域：request每次请求后sesson更新
    System.out.println("token in current session: " + csrf_token_in_session);
    System.out.println("csrf_token: " + csrf_token);
    if (csrf_token == null || !csrf_token.equals(csrf_token_in_session)) {
      response.setStatus(403);
      return "csrf_validate_error!";
    }
    return "form post ok!";
  }

  private static String getMD5(String str) {
    try {
      // 生成一个MD5加密计算摘要
      MessageDigest md = MessageDigest.getInstance("MD5");
      // 计算md5函数
      md.update(str.getBytes());
      // digest()最后确定返回md5 hash值，返回值为8为字符串。因为md5 hash值是16位的hex值，实际上就是8位的字符
      // BigInteger函数则将8位的字符串转换成16位hex值，用字符串来表示；得到字符串形式的hash值
      return new BigInteger(1, md.digest()).toString(16);
    } catch (Exception e) {
      logger.error("MD5加密出现错误", e);
    }
    return str;
  }

  /**
   * 允许浏览器跨域Ajax
   *
   * @param response
   * @author gulixing@msn.com
   * @date 2016年12月28日
   */
  private void addAllowHeader(HttpServletResponse response) {
    response.addHeader("Access-Control-Allow-origin", "*");
    response.addHeader("Access-Control-Allow-Methods", "GET, POST");
  }
}
